#ifndef IO_ENGINE_H
#define IO_ENGINE_H

#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <mutex>
#include <queue>
#include <memory>
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <libaio.h>
#include <cstdint>

#define ANN_BLOCK_SIZE 4096
#define MAX_BATCH_SIZE 64 //IO引擎最多允许一次最多发起多少个读请求，大于这个值将被减小到MAX_BATCH_SIZE
#define SMALL_IO_THRESHOLD 4096
#define IOCB_REQ_PRIO_HIGH 0
#define IOCB_REQ_PRIO_LOW 20

//异步版Manager，支持持续发起请求以及异步获取结果
class IOUringManager_ASYNC {
private:
    int fd_;
    int numRequests_;
    io_context_t ctx_;

public:
    IOUringManager_ASYNC(int fd, int request_count)
        : fd_(fd), numRequests_(request_count) {
        //这里不用检查request_count是否大于MAX_BATCH_SIZE，因为request_count是没有限制的可以提交的时候动态伸缩，MAX_BATCH_SIZE是获取结果的时候才有用
        // if(numRequests_ > MAX_BATCH_SIZE){
        //     std::cerr << "IOUringManager_ASYNC: request_count("<< request_count <<") > MAX_BATCH_SIZE("<< MAX_BATCH_SIZE <<")" << std::endl;
        //     exit(0);
        // }
        memset(&ctx_, 0, sizeof(io_context_t));
        int ret = io_setup(numRequests_, &ctx_);
        if (ret < 0) {
            std::cerr << "Failed to initialize libaio. Error code: " << ret << std::endl;
            exit(EXIT_FAILURE);
        }
    }

    ~IOUringManager_ASYNC() {
        io_destroy(ctx_);
    }

    //ids是每个请求的标识符，由调用者自己定义，将作为getCompletedResultIds的返回值
    int submitRequests(std::vector<void*>& input_bufAddrs,
                      const std::vector<long long>& offsets,
                      const std::vector<long long>& sizes,
                      const std::vector<int>& ids,
                      int input_num_request = 0) {
        size_t num_requests = input_num_request == 0 ? input_bufAddrs.size() : input_num_request;
        if (num_requests == 0) return 0;

        // Validate input vectors
        if (input_bufAddrs.size() < num_requests || offsets.size() < num_requests ||
            sizes.size() < num_requests || ids.size() < num_requests) {
            std::cerr << "Input vectors too small for requested operations" << std::endl;
            return 0;
        }

        std::vector<struct iocb*> iocb_ptrs;
        iocb_ptrs.reserve(num_requests);

        try {
            for (size_t i = 0; i < num_requests; ++i) {
                struct iocb* cb = new struct iocb;
                io_prep_pread(cb, fd_, input_bufAddrs[i], sizes[i], offsets[i]);
                cb->data = (void*)(uintptr_t)(ids[i]); // Store self pointer for cleanup
                //根据IO尺寸设置，如果IO粒度小于SMALL_IO_THRESHOLD，设置为高优先级
                if(num_requests <= 8){
                    cb->aio_reqprio = IOCB_REQ_PRIO_HIGH;
                }
                else{
                    cb->aio_reqprio = IOCB_REQ_PRIO_LOW;
                }
                iocb_ptrs.push_back(cb);
            }
        } catch (...) {
            // Cleanup if memory allocation fails
            for (auto cb : iocb_ptrs) {
                delete cb;
            }
            return 0;
        }

        int ret = io_submit(ctx_, num_requests, iocb_ptrs.data());
        if (ret < 0) {
            // Full cleanup on submission failure
            for (auto cb : iocb_ptrs) {
                delete cb;
            }
            return 0;
        } else if (ret < static_cast<int>(num_requests)) {
            // Partial success cleanup
            int submitted = ret;
            for (size_t i = submitted; i < num_requests; ++i) {
                struct iocb* cb = iocb_ptrs[i];
                delete cb;
            }
            return submitted;
        }
        return ret;
    }

    int getCompletedResultIds(std::vector<int>& completedIds, int min_nr = 0) {
        // struct iocb* cbs[MAX_BATCH_SIZE];
        struct io_event events[MAX_BATCH_SIZE];
        int ret = io_getevents(ctx_, min_nr, MAX_BATCH_SIZE, events, NULL);
        if (ret < 0) {
            std::cerr << "Failed to wait for I/O completion" << std::endl;
            exit(0);
        }

        for (int i = 0; i < ret; ++i) {
            int index = (int)(uintptr_t)events[i].data;
            // cout << "getCompletedResultIndex: index=" << index << endl;
            if (events[i].res < 0) {
                std::cerr << "I/O operation failed for request " << index << std::endl;
            } else {
                completedIds.push_back(index);
            }

            // 获取完成事件对应的iocb对象
            struct iocb* cb = static_cast<struct iocb*>(events[i].obj);
            // 释放iocb内存
            delete cb;
        }
        return ret;
    }
};


//一些便于Direct IO的帮助函数
//计算对齐的块数
int block_count(long long size) {
    return (size + ANN_BLOCK_SIZE - 1) / ANN_BLOCK_SIZE;
}

//计算对齐的大小。输入需要的字节数，返回大于等于需要的字节数的对齐大小.
long long aligned_size(long long size, long long start_addr = 0) {
    return block_count(size + start_addr%ANN_BLOCK_SIZE) * ANN_BLOCK_SIZE;
}

//计算对齐的偏移量，即某个偏移量在对齐IO中，在文件中的起始位置，≤起始地址
long long aligned_offset_in_file(long long offset) {
    return offset / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;
}

//通过aligned_offset_in_file返回的地址中，真实起始地址在分配的buf中的偏移量
long long aligned_offset_in_buf(long long offset) {
    return offset % ANN_BLOCK_SIZE;
}

//给定需要的字节数，返回≥字节数且对齐的内存地址
void* get_aligned_buffer(long long size, long long start_addr = 0) {
    if (size <= 0) {
        std::cout << "[Error] allocate_aligned_buffer size is invalid" << std::endl;
        exit(0);
    }
    // 分配对齐的内存
    void* buf = nullptr;
    if (posix_memalign(&buf, ANN_BLOCK_SIZE, aligned_size(size, start_addr)) != 0) {
        std::cout << "[Error] posix_memalign failed" << std::endl;
        exit(0);
    }
    return buf;
}

//用O_DIRECT方式打开文件后用该方法读取文件
//负责对buf分配空间，并且使其指向真实数据的地址
//现在测试阶段为了方便使用了拷贝，后续进行优化
int read_file_aligned(int fd, char* buf, long long offset, long long size) {
    if(size <= 0) {
        exit(0); // 无效的大小
    }
    // 计算需要的对齐内存大小
    int aligned_size = (size + ANN_BLOCK_SIZE - 1) / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;

    // 分配对齐的内存
    void* aligned_buf = nullptr;
    if (posix_memalign(&aligned_buf, ANN_BLOCK_SIZE, aligned_size) != 0) {
        perror("posix_memalign");
        exit(0);
    }

    // 设置文件偏移量
    int aligned_offset = offset / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;
    if (lseek(fd, aligned_offset, SEEK_SET) == -1) {
        perror("lseek");
        free(aligned_buf);
        exit(0);
    }

    // 读取文件内容到对齐的内存中
    ssize_t bytes_read = read(fd, aligned_buf, aligned_size);
    if (bytes_read == -1 || bytes_read < size) {
        perror("read");
        free(aligned_buf);
        exit(0);
    }

    //计算offset在这个buf中的偏移量
    int offset_in_buf = offset % ANN_BLOCK_SIZE;
    // 将对齐的内存内容复制到buf中
    memcpy(buf, (char*)aligned_buf + offset_in_buf, size);
    return bytes_read;
}

#endif